Repository navigation
[OMNIML-6072] [2/21] Make amax synchronization correct under data-parallel calibration - #2744
shengliangxu wants to merge 11 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 Walkthrough
Merge Risk: 🟡 Moderate · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 51.85% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 27 functions across 5 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Title check The title clearly and concisely describes the main change: correcting amax synchronization for data-parallel calibration. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns PASS. The reviewed range changes three ModelOpt Python modules and no examples or dependency files. Added code only implements distributed amax synchronization and an expert-sharing flag. The added di… Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Full details: Security Anti-Patterns
-
Explanation
PASS. The reviewed range changes three ModelOpt Python modules and no examples or dependency files. Added code only implements distributed amax synchronization and an expert-sharing flag. The added diff contains no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), hardcodedtrust_remote_code=True, external-inputeval()/exec(), or# nosec.SECURITY.mdstates these rules apply to non-test code, and the dependency-file diff is empty.
-
✨ Finishing Touches: 💡 1
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
06a5f2d to
5f115f2
Compare
|
|
/ok to test 5f115f2 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2744 +/- ##
==========================================
- Coverage 79.82% 79.13% -0.69%
==========================================
Files 661 661
Lines 74928 74977 +49
==========================================
- Hits 59810 59332 -478
- Misses 15118 15645 +527
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
5f115f2 to
d6a4b30
Compare
|
/ok to test d6a4b30 |
shengliangxu
changed the title
d6a4b30 to
e28bd09
Compare
e28bd09 to
235ecbd
Compare
|
/ok to test 235ecbd |
shengliangxu
changed the title
235ecbd to
9572dc5
Compare
shengliangxu
changed the title
|
/ok to test 9572dc5 |
9572dc5 to
6e7ca4f
Compare
shengliangxu
changed the title
|
/ok to test 4f9e94b |
|
/ok to test 669b19c |
|
Design suggestion (non-blocking): one sync entry point for calibration #2736 gives the control plane labeled, fail-together collectives. Calibration doesn't get that yet: Suggestion (EP only for now; TP/DP can move behind it later):
Which algorithms are meant to work under EP in this stack? |
|
/ok to test 838be3c |
|
/ok to test 719a36d |
|
@Fridah-nv, thanks. Which algorithms work under expert parallelism, as of this push:
On points 1 and 2, one sync entry point that reads the plan, with labeled, fail-together collectives: agreed on the direction, as a follow-up. It would reshape |
|
/ok to test 481b71b |
|
/ok to test d11f0f8 |
|
/ok to test 7408fe2 |
|
/ok to test ef0ff99 |
With data-parallel calibration a module can be exercised on some ranks and not others -- a sparsely routed MoE expert, or anything a given rank's batch never reaches -- so _amax ends up present on only a subset of the group. Every rank must still join the collective or the participant set mismatches and NCCL hangs, which makes "what does a rank with no amax contribute" the whole problem. For a per-tensor amax the answer is easy: the shape is known locally, so a have-not contributes a scalar zero, the identity for MAX over non-negative amax. If the all-reduced result is zero then nobody calibrated it, and the buffer is dropped so downstream `amax is None` handling still applies -- keeping it as 0 would quantize the module to nothing. A static per-channel or per-block amax has no locally-known shape, so a have-not cannot fabricate a correctly-shaped zero. Ranks agree first, via one MAX all-reduce of a packed presence flag that every rank joins: nobody calibrated it is fine and skipped; everybody did falls through to the real reduce; a genuine split raises on every rank at once. Raising only where the amax is missing would leave the other ranks blocked in a collective that never completes. A dynamic-block quantizer -- NVFP4 activations -- stores only a scalar global scale; its block_sizes drive runtime scales and do not describe the stored amax. Classifying by "has block_sizes" would push it down the static path and turn the ordinary uneven-calibration case into a hard error, so it is treated as per-tensor. Every rank also contributes the same dtype and device. A calibrated amax keeps the calibrator's dtype while a have-not's zero is float32, and NCCL hangs on a mismatch rather than raising (seen on Kimi-K3, NVFP4 EP64: 8 ranks with an amax, 56 without, and the collective never returned), so the amax is cast to float32 first. The device comes from the group's backend, which every rank shares -- CUDA under NCCL, CPU otherwise -- so a gloo group on a CPU host works. max_calibrate keeps main's sync structure -- DP always, EP filtered by _should_sync_amax_across_ep -- so Megatron's numerics are unchanged, with three changes: a rank joins the DP sync even without an amax (disabled quantizers are skipped, which is symmetric since it is a config property); the input amax of expert-parallel experts is shared over the EP group BEFORE the completeness check, so experts that received no tokens pick up the layer's value instead of failing it; and the allocator's cached pool is returned first, since NCCL allocates its buffers outside it. A failed amax all-reduce now says it is most likely that out-of-memory rather than surfacing as a bare CUDA error. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Two reads go through an Optional that is set by then, which the mypy run in pre-commit (and so in CI) rejects: the amax every rank holds once it reaches the MAX all-reduce (calibrated, or the zero synthesised just above), and the parallel state _has_expert_parallelism has already checked. Assert both. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The sync decided that no rank had calibrated a quantizer when the reduced amax was zero, and dropped it. A calibrated amax can legitimately be zero -- an all-zero weight -- and dropping it made the FP8 export fail on the missing _amax for any run with a process group (Megatron, FSDP2, DDP). The ranks now agree through a presence flag carried in the same all-reduce. The sync also promoted every amax to float32 for the collective and left it there, so later bf16 writes (MSE calibration) stopped rounding as before. Each rank now keeps its calibrator's dtype, and a rank that never calibrated the quantizer adopts the group's. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…he release Sharing the routed experts' input amax over the expert-parallel group before the completeness check exists for HF fused experts, whose experts see only the tokens routed to them. It also applied to Megatron's experts, which then stopped failing on incomplete calibration and silently took the group max. Only modules that opt in (_share_input_amax_across_ep, set by the HF fused experts wrapper) share early; Megatron keeps main's check-then-sync. max_calibrate also ran gc.collect() and emptied the CUDA cache on every call, which AWQ and SmoothQuant make once per linear: ~45 s on an 8B model. The distributed forward loop already returns the cache after its last forward, right before the sync this was for. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With one rank nobody lacks the amax, so nothing raises and pytest.raises fails inside the worker. That failure is a BaseException the worker pool does not catch: the worker dies without reporting, and every later test on the session's pool waits on it. CI's PR GPU runner has one GPU. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Each of them spawned two ranks, about 14 s of the linux unit job apiece. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The amax sync is a collective every rank of the group joins, and it handles a rank with no amax of its own. The tensor-parallel callers still skipped such a rank, for linears and for the KV-cache quantizers alike, so a rank that calibrated the quantizer waited in the sync for a peer that never joined it. They now gate on whether the quantizer is enabled, as the data-parallel and expert-parallel sync does: the config enables the same quantizers on every rank, and a disabled one, which nothing quantizes with, joins no collective. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The GPU test file sized its world from the GPU count, so on the 1-GPU PR runner each "some ranks missing" case ran on one rank that held the amax, and its one two-rank case needed two GPUs: no PR exercised the multi-rank behavior. The sync logic runs on gloo through cpu_dist_workers instead: - the uneven per-channel raise, which every rank must reach together, on a group that gives up in seconds rather than hang; - a per-channel quantizer no rank calibrated, which is not an error; - an NVFP4 dynamic-block quantizer, which syncs as a per-tensor one; - the group's max over ranks that all calibrated, which no unit test checked: a rank keeping its own amax went unnoticed. The GPU file keeps the check only NCCL can make, where a rank without an amax puts the zero it contributes, and asks for two GPUs: with one, no rank lacks the amax. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
- test_amax_sync_with_a_rank_that_never_calibrated is the one-calibrated half of test_the_sync_leaves_every_rank_the_calibrators_dtype, which runs the same case in bf16 and checks the dtype too. - test_a_calibrated_zero_amax_survives_the_sync and test_a_quantizer_no_rank_calibrated_stays_uncalibrated become one test_the_sync_keeps_a_calibrated_zero_apart_from_uncalibrated over calibrated-zero, uncalibrated and uncalibrated-per-channel, each case keeping its assert. A have-not rank that keeps no amax or takes zero, a zero amax read as uncalibrated, nobody's amax becoming zero, and a per-channel quantizer nobody calibrated raising each fail a remaining test. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The two jobs that sync on a group with a 10-second timeout left it, with its gloo backend, registered on the session's persistent cpu_dist_workers until the pool shut down; they now destroy it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 6760929 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # rank skips and another syncs pairs mismatched tensors in every later | ||
| # all_reduce, which NCCL hangs on rather than reporting. | ||
| if leaf.is_enabled: | ||
| leaf.sync_amax_across_distributed_group(ep_group) |
There was a problem hiding this comment.
[SUGGESTION] The block comment above states the contract as "sync_amax_across_distributed_group fills a None rank with the group MAX; if NO rank saw tokens the amax is dropped group-wide". That only holds for a per-tensor (or dynamic-block) quantizer. For an input quantizer with axis is not None or a static block_sizes, the new implementation takes the static branch and the cold-expert rank hits the "unevenly-calibrated non-per-tensor quantizer" RuntimeError instead — i.e. exactly the case this loop was added to make work turns into a hard failure, with an error message pointing the user at --calib-size even though no calibration size can fix a cold EP shard.
Shipped MoE activation recipes (FP8 per-tensor, NVFP4 dynamic-block) are all per-tensor-amax, so nothing today is affected; but a per-channel activation recipe on fused experts would regress to an error. Either note the restriction in the comment, or make it explicit where it is cheap to check — e.g. skip-with-warning rather than raise when the quantizer is a routed-expert input quantizer being shared over EP.
There was a problem hiding this comment.
Postponed to a follow-up: the comment should limit the fill to a per-tensor (or dynamic-block) amax. A static per-channel or per-block input amax that only some ranks hold raises on every rank together instead of hanging, and, as you note, no shipped recipe calibrates an expert input that way, so this approved slice keeps the comment as it is.
|
Claude review summary Scope reviewed: full PR (5 files, +358/-15) — Findings: CRITICAL 0, IMPORTANT 0, SUGGESTION 4. Assessment The core design holds up. The three things I checked hardest all came out right:
Test coverage is proportionate and the multi-rank cases run on gloo so every PR exercises them; using a 10 s group timeout so a left-behind rank fails instead of hanging CI is the right pattern per the contributing guide. The GPU file keeps only what needs NCCL. The suggestions None block. In rough order of usefulness:
One out-of-diff note
Risk: low. Behaviour at world size 1 and for Megatron is unchanged, the backward-compatible direction is preserved, and the new failure modes are loud rather than silent. 🤖 Generated with Claude Code |
What does this PR do?
Type of change: Bug fix (for data-parallel calibration, which the series introduces)
With data-parallel calibration a module can be exercised on some ranks and not others: a sparsely routed MoE expert, or anything a given rank's batch never reaches. Its
_amaxthen exists on only part of the group. Every rank must still join the collective, or the participant set mismatches and NCCL hangs. So the whole question is what a rank with no amax contributes.amax is Nonehandling still applies.max_calibratekeeps main's sync structure (data-parallel always, expert-parallel filtered by_should_sync_amax_across_ep), so Megatron's numerics are unchanged. Three changes:_share_input_amax_across_ep) share their input amax over the EP group before the completeness check, so experts that received no tokens take the layer's value instead of failing it. Megatron's experts keep main's check-then-sync, including its "MoE calibration incomplete" error.A failed amax all-reduce now names its likely cause: NCCL running out of memory for its buffers.
Usage
N/A: internal to calibration.
Testing
tests/unit/torch/quantization/test_dist.py, under gloo on CPU, two ranks, so every PR runs the multi-rank cases:Each fails on the code before its fix, or under a mutation of the sync it checks.
tests/gpu/torch/quantization/test_amax_sync_distributed.pykeeps only what needs NCCL: a rank without an amax contributes its zero on the collective's CUDA device, when one rank calibrated and when none did (2 GPUs; the other multi-rank cases moved to the unit test above, since the PR runner has 1 GPU).The GPU results here were measured on earlier revisions of this branch, including the earlier, larger GPU file (7 passed on a 4-GPU node); this revision's GPU tests run in CI only.
Against main, inside a one-rank process group, as Megatron, FSDP2 and DDP runs calibrate: FP8 max, FP8 MSE and INT8 MSE calibration leave every amax value and dtype identical to main's (48 of 48). An FP8 export with an all-zero weight succeeds, as on main.
A calibrated zero amax and the uncalibrated quantizers are now one parametrized test, and the never-calibrated-rank test, a case the dtype test runs, is dropped. Each drop is checked by mutation.
Locally on this revision's content (the series on main 6cd9ced, restacked unchanged onto main 0638a27 after [1/21] merged): pre-commit (every hook, mypy included) over each PR of the stack against main, and
test_dist.pywith the transport's tests on transformers 5.18 and 5.5. CI runs the full unit lanes and the docs build.The series' model runs (see [21/21]) calibrate sparsely routed MoEs at 4 to 64 ranks through this path.
Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part 2 of 21 of expert-parallel PTQ for Hugging Face MoE models. [1/21] is merged, so this PR is now the bottom of the stack, on main.
Summary by CodeRabbit